# -*- coding:utf-8 -*-
"""
    timestamp_low 时间戳
    bfee_count 数据包数量
    Nrx,Ntx 分别表示接收端和发送端的天线数量
    rssi_a, rssi_b, rssi_c 每个天线的RSSI数据，单位dB，
    agc Automatic Gain Control
    perm NIC重排列后的顺序结果，代表RF链路的顺序
    rate 发送包的rate
    csi CSI原始数据，是个Ntx×Nrx×30复数矩阵
"""
from itertools import accumulate

max_antenna = 3
triangle = list(accumulate(range(1, max_antenna + 1)))


class CSI_Data:
    def __init__(self, data):
        self.correct = False
        self.timestamp_low = data[0] + (data[1] << 8) + (data[2] << 16) + (data[3] << 24)
        self.bfee_count = data[4] + (data[5] << 8)
        self.Nrx = data[8]
        self.Ntx = data[9]
        self.rssi_a = data[10]
        self.rssi_b = data[11]
        self.rssi_c = data[12]
        self.noise = data[13]
        if (self.noise > 128):
            self.noise -= 256
        self.agc = data[14]
        antenna_sel = data[15]
        size = data[16] + (data[17] << 8)
        self.rate = data[18] + (data[19] << 8)
        calc_len = (30 * (self.Nrx * self.Ntx * 8 * 2 + 3) + 7) // 8
        if size != calc_len:
            return
        index = 0
        payload = data[20:]
        self.csi = [[[0 for _ in range(30)] for _ in range(self.Nrx)] for _ in range(self.Ntx)]
        for i in range(30):
            index += 3
            remainder = index % 8
            for k in range(self.Nrx):
                for j in range(self.Ntx):
                    tmp = ((payload[index // 8] >> remainder) | (payload[index // 8 + 1] << (8 - remainder))) & 0xFF
                    tmp2 = ((payload[index // 8 + 1] >> remainder) | (
                            payload[index // 8 + 2] << (8 - remainder))) & 0xFF
                    if tmp > 128:
                        tmp -= 256
                    if tmp2 > 128:
                        tmp2 -= 256
                    self.csi[j][k][i] = tmp + tmp2 * 1j
                    index += 16
        self.perm = [0] * 3
        self.perm[0] = (antenna_sel & 0x3) + 1
        self.perm[1] = ((antenna_sel >> 2) & 0x3) + 1
        self.perm[2] = ((antenna_sel >> 4) & 0x3) + 1
        if self.permsum(self.perm,self.Nrx) == triangle[self.Nrx - 1]:
            self.perm = list(map(lambda x: x - 1, self.perm))
            self.changeperm(self.perm)
            self.correct = True

    def changeperm(self, perm):
        csinew = [[[0 for _ in range(30)] for _ in range(self.Nrx)] for _ in range(self.Ntx)]
        for i in range(30):
            for j in range(self.Ntx):
                for k in range(self.Nrx):
                    csinew[j][k][i] = self.csi[j][perm[k]][i]
        self.csi = csinew


    def permsum(self,perm,rxnum):
        ret = 0
        for i in range(rxnum):
            ret += perm[i]
        return ret